package at.ofai.gate.riplugin2;

import gate.util.GateRuntimeException;
import java.io.File;
import java.io.IOException;
import java.util.LinkedList;
import java.util.List;
import org.apache.lucene.index.CorruptIndexException;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.Term;
import org.apache.lucene.index.TermEnum;
import org.apache.lucene.misc.HighFreqTerms;
import org.apache.lucene.store.FSDirectory;
import org.apache.lucene.util.PriorityQueue;

/**
 * Class to return the N most frequent terms for a Lucene index. Has a
 * Main function to print those terms out to standard output.
 * The main function can also be used to print out all the terms and their
 * document frequencies if 0 is given for N.
 * <p>
 * Parts of this class are based on org.apache.lucene.misc.HighFreqTerms
 * 
 * @author Johann Petrak
 */
public class LuceneTermsInfo {

  public List<TermInfo> getMostFrequentTerms(
    IndexReader reader, String fieldName, String measure, int n, int numDocs) throws IOException, Exception {
    LinkedList<TermInfo> ret = new LinkedList<TermInfo>();
    TermInfoQueue tiq = new TermInfoQueue(measure,n);
    TermEnum terms = reader.terms();
    while (terms.next()) {
      if (terms.term().field().equals(fieldName)) {
        int df = terms.docFreq();
        double idf = Math.log(numDocs/(double)(df+1))+1.0;
        long tf = HighFreqTerms.getTotalTermFreq(reader, terms.term());
        double tfidf = (double)tf * idf;
        tiq.insertWithOverflow(new TermInfo(terms.term(), df, idf, tf, tfidf));
      }
    }
    while (tiq.size() != 0) {
      TermInfo termInfo = tiq.pop();
      ret.addFirst(termInfo);
    }
    return ret;
  }
  
  
  // TODO: add more functions:
  // get terms by term frequency, document frequency, idf, tf/idf etc.
  // show version information
  // show field names
  // show documents that contain a term
  // ...
  // Collection<String> indexReader.getFieldNames(IndexReader.FieldOption.ALL)
  // long indexReader.getUniqueTermCount() 
  // int numDocs()
  // TODO: in newer versions of Lucene getting the N most frequent (by DF)
  // terms can be done usingthe 
  // HighFreqTerms.getHighFreqTerms(IndexReader, int numTerms, String fieldname) 
  // method.
  // Look at sortByTotalTermFreq
  public static void main(String args[]) throws CorruptIndexException, IOException, Exception {
    if (args.length !=2 && args.length != 3 && args.length != 4) {
      usage();
      System.exit(1);
    }
    String dirname = args[0];
    String sortby = args[1];
    if(!sortby.equals("numdocs") && args.length < 3) {
      usage();
      System.exit(1);
    }
    IndexReader indexreader =
      IndexReader.open(FSDirectory.open(new File(dirname)), true);
    int numDocs = indexreader.numDocs();    
    int n;
    if(args.length == 3) {
      n = Integer.parseInt(args[2]);
    } else {
      System.err.println(numDocs);
      return;
    }
    String field = "contents";
    if(args.length == 4) {
      field = args[3];
    }
    if(n == 0) {
      TermEnum terms = indexreader.terms();
       while (terms.next()) {
         if (terms.term().field().equals(field)) {
           int df = terms.docFreq();
           double idf = Math.log((double)numDocs/(double)(df+1))+1.0;
           long tf = HighFreqTerms.getTotalTermFreq(indexreader, terms.term());
           double tfidf = (double)tf * idf;
           System.out.println(terms.term().text()+ "\t"+ 
             df + "\t" + idf + "\t" + tf + "\t" + tfidf);
         }
       }
    } else {
      LuceneTermsInfo lmft = new LuceneTermsInfo();
      List<TermInfo> tis = lmft.getMostFrequentTerms(indexreader, field, sortby, n, numDocs);
      for (TermInfo ti : tis) {
        System.out.println(ti.term.text() + "\t" + ti.df +
          "\t" + ti.idf + "\t" + ti.tf + "\t" + ti.tfidf);
      }
    }
  }

  private static void usage() {
      System.err.println("Need either 3 or 4 args: indexdir df|idf|tf|tfidf nterms [fieldname] (if nterms=0, all terms, unsorted)");
      System.err.println("Or 2 args: indexdir numdocs");    
  }
  
  public final class TermInfo {

    public TermInfo(Term t, int df, double idf, long tf, double tfidf) {
      term = t;
      this.df = df;
      this.idf = idf;
      this.tf = tf;
      this.tfidf = tfidf;
    }
    public int df;
    public long tf;
    public double idf;
    public double tfidf;
    public Term term;
  }

  final class TermInfoQueue extends PriorityQueue<TermInfo> {

    TermInfoQueue(String m, int size) {
      if(m.equals("df") || m.equals("tf") || m.equals("idf") ||
         m.equals("tfidf")) {
        this.m = m;
      } else {
        throw new GateRuntimeException("Invalid measure: "+m);
      }
      initialize(size);
    }
    String m;
    @Override
    protected final boolean lessThan(TermInfo termInfoA, TermInfo termInfoB) {
      boolean ret;
      if(m.equals("df")) {
        return termInfoA.df < termInfoB.df;
      } else if(m.equals("tf")) {
        return termInfoA.tf < termInfoB.tf;
      } else if(m.equals("idf")) {
        return termInfoA.idf > termInfoB.idf;
      } else if(m.equals("tfidf")) {
        return termInfoA.tfidf < termInfoB.tfidf;
      } else {
        // should never happen, but makes compiler happy!
        return false;
      }
    }
  }
}
